#!/usr/bin/env python3

"""
Complete Neural-SLAM Enhancement Pipeline
One-click script to run the entire workflow from setup to paper generation
"""

import subprocess
import sys
import os
import time
import json
from datetime import datetime

class NeuralSLAMPipeline:
    """Complete pipeline manager for Neural-SLAM enhancement project"""
    
    def __init__(self, experiment_name=None):
        self.start_time = time.time()
        self.experiment_name = experiment_name or f"neural_slam_enhanced_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        self.status = {}
        
        print("🚀 Neural-SLAM Complete Enhancement Pipeline")
        print("=" * 60)
        print(f"📊 Experiment: {self.experiment_name}")
        print(f"🕒 Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print("=" * 60)
    
    def run_phase(self, phase_name, description, command, timeout=None):
        """Run a pipeline phase with error handling"""
        
        print(f"\n📅 Phase: {description}")
        print("-" * 50)
        print(f"🏃 Command: {command}")
        
        start_time = time.time()
        
        try:
            if isinstance(command, str):
                # Shell command
                result = subprocess.run(command, shell=True, capture_output=True, 
                                      text=True, timeout=timeout)
            else:
                # Python function
                result = command()
                if result is None:
                    result = type('Result', (), {'returncode': 0, 'stdout': '', 'stderr': ''})()
            
            duration = time.time() - start_time
            
            if hasattr(result, 'returncode') and result.returncode == 0:
                print(f"✅ {description} completed successfully in {duration:.1f}s")
                self.status[phase_name] = {
                    'status': 'success',
                    'duration': duration,
                    'message': 'Completed successfully'
                }
                return True
            else:
                error_msg = getattr(result, 'stderr', 'Unknown error')
                print(f"❌ {description} failed: {error_msg}")
                self.status[phase_name] = {
                    'status': 'failed',
                    'duration': duration,
                    'error': error_msg
                }
                return False
                
        except subprocess.TimeoutExpired:
            print(f"⏰ {description} timed out after {timeout}s")
            self.status[phase_name] = {
                'status': 'timeout',
                'duration': timeout,
                'error': 'Timeout'
            }
            return False
            
        except Exception as e:
            duration = time.time() - start_time
            print(f"❌ {description} failed with exception: {str(e)}")
            self.status[phase_name] = {
                'status': 'error',
                'duration': duration,
                'error': str(e)
            }
            return False
    
    def phase_1_environment_setup(self):
        """Phase 1: Environment setup and validation"""
        
        def setup_environment():
            """Setup and test environment"""
            
            # Check if we're in the right directory
            if not os.path.exists('main.py'):
                raise Exception("Not in Neural-SLAM directory. Please run from the project root.")
            
            # Create necessary directories
            os.makedirs('./tmp', exist_ok=True)
            os.makedirs('./logs', exist_ok=True)
            os.makedirs('./experiments', exist_ok=True)
            os.makedirs('./paper_results', exist_ok=True)
            
            print("📁 Created necessary directories")
            
            # Test environment
            import test_env
            success = test_env.main()
            
            if not success:
                raise Exception("Environment validation failed")
            
            return type('Result', (), {'returncode': 0})()
        
        return self.run_phase(
            'env_setup',
            'Environment Setup and Validation',
            setup_environment,
            timeout=300
        )
    
    def phase_2_baseline_extraction(self):
        """Phase 2: Extract baseline metrics"""
        
        def extract_baseline():
            """Extract baseline performance metrics"""
            
            import extract_baseline_metrics
            report = extract_baseline_metrics.main()
            
            print("📊 Baseline metrics extracted successfully")
            return type('Result', (), {'returncode': 0})()
        
        return self.run_phase(
            'baseline_extraction',
            'Baseline Metrics Extraction',
            extract_baseline,
            timeout=1800  # 30分钟，因为可能需要训练
        )
    
    def phase_3_model_improvements(self):
        """Phase 3: Apply model improvements"""
        
        def apply_improvements():
            """Apply model improvements"""
            
            # Backup original model
            if os.path.exists('model.py'):
                import shutil
                shutil.copy('model.py', 'model_original_backup.py')
                print("📦 Backed up original model.py")
            
            # Copy improved model
            import shutil
            shutil.copy('model_improved.py', 'model.py')
            print("✨ Applied improved model with attention mechanisms")
            
            return type('Result', (), {'returncode': 0})()
        
        return self.run_phase(
            'model_improvements',
            'Apply Model Improvements',
            apply_improvements,
            timeout=60
        )
    
    def phase_4_quick_training(self):
        """Phase 4: Quick training test"""
        
        cmd = f"""python main.py \
            --exp_name {self.experiment_name}_quick_test \
            --num_episodes 100 \
            --num_episode_length 50 \
            --train_slam 1 \
            --train_global 1 \
            --train_local 1 \
            --save_interval 50 \
            --log_interval 10 \
            --print_images 0"""
        
        return self.run_phase(
            'quick_training',
            'Quick Training Test (100 episodes)',
            cmd,
            timeout=1800  # 30 minutes
        )
    
    def phase_5_full_training(self):
        """Phase 5: Full training run"""
        
        cmd = f"""python training_pipeline.py"""
        
        return self.run_phase(
            'full_training',
            'Full Training Pipeline',
            cmd,
            timeout=43200  # 12 hours
        )
    
    def phase_6_evaluation(self):
        """Phase 6: Model evaluation"""
        
        cmd = f"""python main.py \
            --eval 1 \
            --split val \
            --num_episodes 50 \
            --exp_name {self.experiment_name}_final_eval \
            --print_images 1 \
            --save_trajectory_data 1"""
        
        return self.run_phase(
            'evaluation',
            'Final Model Evaluation',
            cmd,
            timeout=3600  # 1 hour
        )
    
    def phase_7_paper_generation(self):
        """Phase 7: Generate paper results"""
        
        def generate_results():
            """Generate paper results"""
            
            import generate_paper_results
            generate_paper_results.main()
            
            print("📄 Paper results generated successfully")
            return type('Result', (), {'returncode': 0})()
        
        return self.run_phase(
            'paper_generation',
            'Paper Results Generation',
            generate_results,
            timeout=300
        )
    
    def generate_final_report(self):
        """Generate final pipeline report"""
        
        total_duration = time.time() - self.start_time
        
        report = {
            'experiment_name': self.experiment_name,
            'start_time': datetime.fromtimestamp(self.start_time).isoformat(),
            'end_time': datetime.now().isoformat(),
            'total_duration_hours': total_duration / 3600,
            'phases': self.status
        }
        
        # Count successes
        successful_phases = sum(1 for phase in self.status.values() if phase['status'] == 'success')
        total_phases = len(self.status)
        
        print("\n" + "=" * 60)
        print("🎯 NEURAL-SLAM ENHANCEMENT PIPELINE REPORT")
        print("=" * 60)
        print(f"📊 Experiment: {self.experiment_name}")
        print(f"⏱️  Total Duration: {total_duration/3600:.1f} hours")
        print(f"✅ Success Rate: {successful_phases}/{total_phases} phases")
        print("\n📋 Phase Results:")
        
        for phase_name, phase_data in self.status.items():
            status_icon = {
                'success': '✅',
                'failed': '❌', 
                'timeout': '⏰',
                'error': '💥'
            }.get(phase_data['status'], '❓')
            
            duration = phase_data['duration']
            print(f"   {status_icon} {phase_name}: {phase_data['status']} ({duration:.1f}s)")
        
        print("\n📁 Generated Files:")
        
        # List key output files
        output_files = [
            './paper_results/',
            './experiments/',
            './logs/',
            'baseline_performance.json',
            'model_improved.py'
        ]
        
        for file_path in output_files:
            if os.path.exists(file_path):
                print(f"   ✅ {file_path}")
            else:
                print(f"   ❌ {file_path} (missing)")
        
        print("=" * 60)
        
        if successful_phases == total_phases:
            print("🎉 PIPELINE COMPLETED SUCCESSFULLY!")
            print("📄 Ready for paper submission preparation.")
        else:
            print("⚠️  Pipeline completed with some failures.")
            print("🔧 Please check the failed phases and retry if needed.")
        
        print("=" * 60)
        
        # Save report
        report_file = f'pipeline_report_{self.experiment_name}.json'
        with open(report_file, 'w') as f:
            json.dump(report, f, indent=2)
        
        print(f"📊 Full report saved to: {report_file}")
        
        return report

def main():
    """Run the complete pipeline"""
    
    # Parse command line arguments
    import argparse
    parser = argparse.ArgumentParser(description='Neural-SLAM Complete Enhancement Pipeline')
    parser.add_argument('--experiment-name', type=str, help='Experiment name')
    parser.add_argument('--skip-training', action='store_true', help='Skip long training phases')
    parser.add_argument('--quick-mode', action='store_true', help='Run in quick mode (shorter phases)')
    parser.add_argument('--phase', type=str, help='Run specific phase only')
    
    args = parser.parse_args()
    
    # Initialize pipeline
    pipeline = NeuralSLAMPipeline(args.experiment_name)
    
    # Define phases
    phases = [
        ('phase_1', pipeline.phase_1_environment_setup),
        ('phase_2', pipeline.phase_2_baseline_extraction),
        ('phase_3', pipeline.phase_3_model_improvements),
        ('phase_4', pipeline.phase_4_quick_training),
        ('phase_5', pipeline.phase_5_full_training),
        ('phase_6', pipeline.phase_6_evaluation),
        ('phase_7', pipeline.phase_7_paper_generation)
    ]
    
    # Run specific phase if requested
    if args.phase:
        phase_func = getattr(pipeline, f'phase_{args.phase}', None)
        if phase_func:
            success = phase_func()
            if success:
                print(f"✅ Phase {args.phase} completed successfully")
            else:
                print(f"❌ Phase {args.phase} failed")
                sys.exit(1)
        else:
            print(f"❌ Unknown phase: {args.phase}")
            sys.exit(1)
        return
    
    try:
        # Run all phases
        for phase_name, phase_func in phases:
            
            # Skip training phases if requested
            if args.skip_training and 'training' in phase_name:
                print(f"⏭️  Skipping {phase_name} (--skip-training)")
                continue
            
            # Run phase
            success = phase_func()
            
            if not success:
                print(f"\n❌ Pipeline stopped at {phase_name}")
                if input("Continue anyway? (y/N): ").lower() != 'y':
                    break
            
            # Quick mode: shorter phases
            if args.quick_mode and phase_name == 'phase_4':
                print("⚡ Quick mode: skipping remaining training phases")
                break
        
    except KeyboardInterrupt:
        print("\n⚠️ Pipeline interrupted by user")
    
    except Exception as e:
        print(f"\n💥 Pipeline failed with exception: {str(e)}")
        raise
    
    finally:
        # Generate final report
        pipeline.generate_final_report()

if __name__ == "__main__":
    main()
